﻿using Microsoft.Win32;
using System.Collections.Generic;
using System.Runtime.InteropServices;

namespace System.Diagnostics
{
	public enum ProcessIntegrityLevel
	{
		Untrusted,
		Undefined,
		Low,
		Medium,
		High,
		System
	}

	internal static class ProcessExtension
	{
		/// <summary> 
		/// The function gets the elevation information of the current process. It  
		/// dictates whether the process is elevated or not. Token elevation is only  
		/// available on Windows Vista and newer operating systems, thus  
		/// IsProcessElevated throws a C++ exception if it is called on systems prior  
		/// to Windows Vista. It is not appropriate to use this function to determine  
		/// whether a process is run as administartor. 
		/// </summary> 
		/// <returns> 
		/// Returns true if the process is elevated. Returns false if it is not. 
		/// </returns> 
		/// <exception cref="System.ComponentModel.Win32Exception"> 
		/// When any native Windows API call fails, the function throws a Win32Exception  
		/// with the last error code. 
		/// </exception> 
		/// <exception cref="System.ArgumentNullException"><paramref name="p"/> must be a valid <see cref="Process"/>.</exception>
		/// <remarks> 
		/// TOKEN_INFORMATION_CLASS provides TokenElevationType to check the elevation  
		/// type (TokenElevationTypeDefault / TokenElevationTypeLimited /  
		/// TokenElevationTypeFull) of the process. It is different from TokenElevation  
		/// in that, when UAC is turned off, elevation type always returns  
		/// TokenElevationTypeDefault even though the process is elevated (Integrity  
		/// Level == High). In other words, it is not safe to say if the process is  
		/// elevated based on elevation type. Instead, we should use TokenElevation.  
		/// </remarks> 
		public static bool IsElevated(this Process p)
		{
			if (p == null)
				throw new ArgumentNullException(nameof(p));

			bool fIsElevated = false;

			// Open the access token of the current process with TOKEN_QUERY. 
			NativeMethods.SafeTokenHandle hToken = NativeMethods.SafeTokenHandle.FromProcess(p.Handle, NativeMethods.AccessTypes.TokenQuery | NativeMethods.AccessTypes.TokenDuplicate);

			// Marshal the TOKEN_ELEVATION struct from native to .NET object. 
			NativeMethods.TOKEN_ELEVATION elevation = hToken.GetInfo<NativeMethods.TOKEN_ELEVATION>(NativeMethods.TOKEN_INFORMATION_CLASS.TokenElevation);

			// TOKEN_ELEVATION.TokenIsElevated is a non-zero value if the token  
			// has elevated privileges; otherwise, a zero value. 
			fIsElevated = (elevation.TokenIsElevated != 0);

			return fIsElevated;
		}

		/// <summary> 
		/// The function gets the integrity level of the current process. Integrity  
		/// level is only available on Windows Vista and newer operating systems, thus  
		/// GetProcessIntegrityLevel throws an exception if it is called on systems  
		/// prior to Windows Vista. 
		/// </summary> 
		/// <returns> 
		/// Returns the integrity level of the current process.
		/// </returns> 
		/// <exception cref="System.ComponentModel.Win32Exception"> 
		/// When any native Windows API call fails, the function throws a Win32Exception  
		/// with the last error code. 
		/// </exception> 
		/// <exception cref="System.ArgumentNullException"><paramref name="p"/> must be a valid <see cref="Process"/>.</exception>
		public static ProcessIntegrityLevel GetIntegrityLevel(this Process p)
		{
			if (p == null)
				throw new ArgumentNullException(nameof(p));

			// Open the access token of the current process with TOKEN_QUERY. 
			NativeMethods.SafeTokenHandle hToken = NativeMethods.SafeTokenHandle.FromProcess(p.Handle, NativeMethods.AccessTypes.TokenQuery | NativeMethods.AccessTypes.TokenDuplicate);

			// Marshal the TOKEN_MANDATORY_LABEL struct from native to .NET object. 
			NativeMethods.TOKEN_MANDATORY_LABEL tokenIL = hToken.GetInfo<NativeMethods.TOKEN_MANDATORY_LABEL>(NativeMethods.TOKEN_INFORMATION_CLASS.TokenIntegrityLevel);

			// Integrity Level SIDs are in the form of S-1-16-0xXXXX. (e.g.  
			// S-1-16-0x1000 stands for low integrity level SID). There is one  
			// and only one subauthority. 
			IntPtr pIL = NativeMethods.GetSidSubAuthority(tokenIL.Label.Sid, 0);
			int il = Marshal.ReadInt32(pIL);
			if (il == 0)
				return ProcessIntegrityLevel.Untrusted;
			if (il == 0x1000)
				return ProcessIntegrityLevel.Low;
			if (il >= 0x2000 && il < 0x3000)
				return ProcessIntegrityLevel.Medium;
			if (il >= 0x4000)
				return ProcessIntegrityLevel.System;
			if (il >= 0x3000)
				return ProcessIntegrityLevel.High;
			return ProcessIntegrityLevel.Undefined;
		}

		/// <summary>
		/// Gets the parent process.
		/// </summary>
		/// <returns>A <see cref="Process"/> object for the process that called the specified process. <c>null</c> if no parent can be established.</returns>
		/// <exception cref="System.ArgumentNullException"><paramref name="p"/> must be a valid <see cref="Process"/>.</exception>
		public static Process GetParentProcess(this Process p)
		{
			if (p == null)
				throw new ArgumentNullException(nameof(p));
			try
			{
				var mo = new System.Management.ManagementObject("win32_process.handle='" + p.Id + "'");
				mo.Get();
				return Process.GetProcessById(Convert.ToInt32(mo["ParentProcessId"]), p.MachineName);
			}
			catch { }
			return null;
		}

		/// <summary>
		/// Gets the child processes.
		/// </summary>
		/// <returns>A <see cref="IEnumerable{Process}"/> reference for enumerating child processes.</returns>
		/// <exception cref="System.ArgumentNullException"><paramref name="p"/> must be a valid <see cref="Process"/>.</exception>
		public static IEnumerable<Process> GetChildProcesses(this Process p)
		{
			if (p == null)
				throw new ArgumentNullException(nameof(p));
			return GetChildProcesses(p, Process.GetProcesses(p.MachineName));
		}

		private static IEnumerable<Process> GetChildProcesses(Process p, Process[] allProcs, bool allChildren = true)
		{
			foreach (Process val in allProcs)
			{
				if (val.Id == p.Id)
				{
					if (allChildren)
						foreach (Process cval in GetChildProcesses(val, allProcs, allChildren))
							yield return cval;
					yield return val;
				}
			}
		}
	}
}
